#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
北京交通大学系统登录模块
负责MIS系统登录和AA系统访问
"""

import pickle
import re
import time
from urllib.parse import urlencode

from dotenv import load_dotenv

load_dotenv()

import requests
from langchain_core.messages import HumanMessage
# 添加Langchain相关导入
from langchain_openai import ChatOpenAI


class BJTSULogin:
    def __init__(self):
        self.session = requests.Session()
        # 设置浏览器请求头
        self.session.headers.update({
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/141.0.0.0 Safari/537.36 Edg/141.0.0.0',
            'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7',
            'Accept-Language': 'zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6',
            'Accept-Encoding': 'gzip, deflate, br, zstd',
            'Connection': 'keep-alive',
            'Upgrade-Insecure-Requests': '1',
            'Cache-Control': 'no-cache',
            'Pragma': 'no-cache',
        })

        # 智谱AI API密钥
        # 初始化Langchain的ChatOpenAI模型
        self.llm = ChatOpenAI(
            model="GLM-4.1V-Thinking-Flash",  # 使用视觉模型
            api_key="f1083622e2424058bfb7cdf7e48a7fb4.XxRmRC4VFmbNjz0K",
            base_url="https://open.bigmodel.cn/api/paas/v4"
        )

    def get_mis_login_url(self):
        """生成MIS系统登录URL"""
        state = str(int(time.time()))
        params = {
            'next': '/o/authorize/',
            'response_type': 'code',
            'client_id': 'aGex8GLTLueDZ0nW2tD3DwXnSA3F9xeFimirvhfo',
            'redirect_uri': 'https://mis.bjtu.edu.cn/auth/callback/?redirect_to=/home/',
            'state': state,
        }
        query_string = urlencode(params, safe='/:?=&')
        return f"https://cas.bjtu.edu.cn/auth/login/?{query_string}", state

    def recognize_captcha_with_vision_model(self, captcha_url):
        """使用视觉模型识别验证码"""
        try:

            # 使用Langchain基于OPENAI格式调用
            if not self.llm:
                print("[错误] 未正确初始化视觉模型")
                return None

            # 构造包含图像URL的消息
            message = HumanMessage(
                content=[
                    {"type": "text", "text": "请识别这张图片中的验证码，需要计算，只返回数字，不要有任何其他文字"},
                    {"type": "image_url", "image_url": {"url": captcha_url}},
                ]
            )

            response = self.llm.invoke([message])

            captcha_text = response.content.strip()
            # 确保只返回数字
            captcha_text = re.sub(r'\D', '', captcha_text)

            print(f"[成功] 视觉模型识别验证码成功: {captcha_text}")
            return captcha_text

        except ImportError:
            print("[错误] 未安装必要的依赖，请执行: pip install langchain langchain-openai")
            return None
        except Exception as e:
            print(f"[错误] 视觉模型识别验证码出错: {e}")
            return None

    def login_mis_system(self, username, password):
        """登录MIS系统"""
        print("正在登录MIS系统...")

        login_url, state = self.get_mis_login_url()

        # 获取登录页面
        try:
            response = self.session.get(login_url)
            if response.status_code != 200:
                print(f"[错误] 无法访问登录页面: {response.status_code}")
                return False

            # 提取CSRF token和验证码
            csrf_pattern = r'name="csrfmiddlewaretoken" value="([^"]+)"'
            csrf_match = re.search(csrf_pattern, response.text)

            yzm_pattern = r'<div class="yzm">(.*?)</div>'
            yzm_match = re.search(yzm_pattern, response.text, re.DOTALL)

            if not csrf_match or not yzm_match:
                print("[错误] 无法提取登录表单信息")
                return False

            csrf_token = csrf_match.group(1)

            # 提取验证码信息
            img_pattern = r'<img src="(/image/[^"]+)"'
            captcha_0_pattern = r'name="captcha_0" value="([^"]+)"'

            img_match = re.search(img_pattern, yzm_match.group(1))
            captcha_0_match = re.search(captcha_0_pattern, yzm_match.group(1))

            if not img_match or not captcha_0_match:
                print("[错误] 无法提取验证码信息")
                return False

            img_path = img_match.group(1)
            captcha_0_value = captcha_0_match.group(1)
            captcha_url = f"https://cas.bjtu.edu.cn{img_path}"

            # 尝试使用视觉模型自动识别验证码
            captcha_text = self.recognize_captcha_with_vision_model(captcha_url)

            # 如果视觉模型识别失败，则保存图片并手动输入
            if not captcha_text:
                # 获取验证码图片
                captcha_response = self.session.get(captcha_url)
                with open("../captcha.png", "wb") as f:
                    f.write(captcha_response.content)
                print("[成功] 验证码已保存为 captcha.png")
                captcha_text = input("请输入验证码: ").strip()

            # 执行登录
            oauth_params = {
                'response_type': 'code',
                'client_id': 'aGex8GLTLueDZ0nW2tD3DwXnSA3F9xeFimirvhfo',
                'redirect_uri': 'https://mis.bjtu.edu.cn/auth/callback/?redirect_to=/home/',
                'state': state,
            }

            next_param = '/o/authorize/?' + urlencode(oauth_params, safe='/:?=&')

            login_data = {
                'next': next_param,
                'csrfmiddlewaretoken': csrf_token,
                'loginname': username,
                'password': password,
                'captcha_0': captcha_0_value,
                'captcha_1': captcha_text,
            }

            login_headers = {
                'Origin': 'https://cas.bjtu.edu.cn',
                'Referer': login_url,
                'Content-Type': 'application/x-www-form-urlencoded',
            }

            response = self.session.post(login_url, data=login_data, headers=login_headers, allow_redirects=False)

            if response.status_code == 302:
                # 处理OAuth2重定向
                authorize_url = response.headers.get('Location', '')

                if '/o/authorize/' in authorize_url:
                    if not authorize_url.startswith('http'):
                        authorize_url = f"https://cas.bjtu.edu.cn{authorize_url}"

                    auth_response = self.session.get(authorize_url, allow_redirects=False)

                    if auth_response.status_code == 302:
                        callback_url = auth_response.headers.get('Location', '')
                        # 处理MIS回调
                        return self.handle_mis_callback(callback_url)
                    else:
                        print(f"[错误] OAuth2授权失败")
                else:
                    print(f"[错误] 重定向URL不包含授权信息")
            else:
                print(f"[错误] 登录失败")

                # 检查是否有错误信息
                if "验证码" in response.text or "captcha" in response.text.lower():
                    print("[警告]  可能是验证码错误")
                elif "密码" in response.text or "password" in response.text.lower():
                    print("[警告]  可能是密码错误")
                elif "用户名" in response.text or "username" in response.text.lower():
                    print("[警告]  可能是用户名错误")

            return False

        except Exception as e:
            print(f"[错误] 登录过程出错: {e}")
            return False

    def handle_mis_callback(self, callback_url):
        """处理MIS系统回调"""
        print("正在处理MIS系统回调...")

        try:
            # 清理冲突的cookies
            cookies_to_remove = []
            for cookie in self.session.cookies:
                if cookie.name in ['csrftoken', 'sessionid'] and cookie.domain != 'mis.bjtu.edu.cn':
                    cookies_to_remove.append(cookie)

            for cookie in cookies_to_remove:
                self.session.cookies.clear(cookie.domain, cookie.path, cookie.name)

            # 发送回调请求
            callback_headers = {
                'Referer': 'https://cas.bjtu.edu.cn/',
                'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7',
            }

            response = self.session.get(callback_url, allow_redirects=False, headers=callback_headers)

            if response.status_code == 302:
                redirect_url = response.headers.get('Location', '')
                if redirect_url.startswith('/'):
                    redirect_url = f"https://mis.bjtu.edu.cn{redirect_url}"

                # 跟随重定向到MIS首页
                mis_headers = {
                    'Referer': callback_url,
                    'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7',
                }

                final_response = self.session.get(redirect_url, headers=mis_headers, allow_redirects=True)

                if final_response.status_code == 200 and "mis.bjtu.edu.cn" in final_response.url:
                    if "logout" in final_response.text.lower() or "退出" in final_response.text:
                        print("[成功] MIS系统登录成功")
                        return True

            print("[错误] MIS系统回调失败")
            return False

        except Exception as e:
            print(f"[错误] 处理MIS回调出错: {e}")
            return False

    def access_aa_system(self):
        """通过MIS系统访问AA系统"""
        print("正在访问AA系统...")

        try:
            # 访问MIS系统的AA模块
            module_url = "https://mis.bjtu.edu.cn/module/module/10/"

            module_headers = {
                'Referer': 'https://mis.bjtu.edu.cn/home/',
                'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7',
            }

            module_response = self.session.get(module_url, headers=module_headers, allow_redirects=True)

            # 检查是否重定向到AA系统
            if "aa.bjtu.edu.cn" in module_response.url:
                print("[成功] 成功访问AA系统")
            else:
                # 手动处理重定向表单
                if module_response.status_code == 200:
                    form_pattern = r'<form[^>]*action=["\']([^"\']*)["\'][^>]*id=["\']redirect["\'][^>]*>'
                    form_match = re.search(form_pattern, module_response.text, re.IGNORECASE)

                    if form_match:
                        action_url = form_match.group(1)
                        # 手动执行重定向
                        module_response = self.session.get(action_url, allow_redirects=True)

            if "aa.bjtu.edu.cn" in module_response.url:
                print("[成功] AA系统访问成功")
                return True

            print("[错误] 未能访问AA系统")
            return False

        except Exception as e:
            print(f"[错误] 访问AA系统出错: {e}")
            return False

    def save_session(self, filename='session.pkl'):
        """保存会话到文件"""
        try:
            with open(filename, 'wb') as f:
                pickle.dump(self.session.cookies, f)
            print(f"[成功] 会话已保存到 {filename}")
            return True
        except Exception as e:
            print(f"[错误] 保存会话失败: {e}")
            return False

    def load_session(self, filename='session.pkl'):
        """从文件加载会话"""
        try:
            with open(filename, 'rb') as f:
                cookies = pickle.load(f)
                self.session.cookies.update(cookies)
            print(f"[成功] 已从 {filename} 加载会话")
            return True
        except FileNotFoundError:
            print(f"[错误] 未找到会话文件 {filename}")
            return False
        except Exception as e:
            print(f"[错误] 加载会话失败: {e}")
            return False

    def get_session(self):
        """获取当前会话"""
        return self.session
